﻿using System;
using System.Diagnostics;
using System.IO;
using System.ServiceModel;
using System.ServiceModel.Channels;
using System.ServiceModel.Dispatcher;
using System.Text;
using System.Xml;

namespace RockBus.ServiceModel.Diagnostics
{
    [Flags]
    public enum MessageLoggingStagesType
    {
        None = 0,
        AfterReceiveRequest = 1,
        BeforeSendReply = 2,
        AfterReceiveReply = 4,
        BeforeSendRequest = 8
    }

    public class MessageLoggingMessageInspector : IDispatchMessageInspector, IClientMessageInspector
    {
        public MessageLoggingMessageInspector(MessageLoggingStagesType messageLoggingStages)
        {
            this._messageLoggingStages = messageLoggingStages;
        }

        public MessageLoggingMessageInspector(MessageLoggingStagesType messageLoggingStages, int maxBufferSize)
            : this(messageLoggingStages)
        {
            this._maxBufferSize = maxBufferSize;
        }

        private readonly MessageLoggingStagesType _messageLoggingStages;

        public MessageLoggingStagesType MessageLoggingStages
        {
            get { return this._messageLoggingStages; }
        }

        private readonly int _maxBufferSize = 65536;

        public int MaxBufferSize
        {
            get { return this._maxBufferSize; }
        }

        #region IDispatchMessageInspector Members

        public object AfterReceiveRequest(ref Message request, IClientChannel channel, InstanceContext instanceContext)
        {
            this.LogMessage(MessageLoggingStagesType.AfterReceiveRequest, ref request);
            return null;
        }

        public void BeforeSendReply(ref Message reply, object correlationState)
        {
            this.LogMessage(MessageLoggingStagesType.BeforeSendReply, ref reply);
        }

        #endregion IDispatchMessageInspector Members

        #region IClientMessageInspector Members

        public void AfterReceiveReply(ref Message reply, object correlationState)
        {
            this.LogMessage(MessageLoggingStagesType.AfterReceiveReply, ref reply);
        }

        public object BeforeSendRequest(ref Message request, IClientChannel channel)
        {
            this.LogMessage(MessageLoggingStagesType.BeforeSendRequest, ref request);
            return null;
        }

        #endregion IClientMessageInspector Members

        private void LogMessage(MessageLoggingStagesType messageLoggingStage, ref Message message)
        {
            if ((this.MessageLoggingStages & messageLoggingStage) != messageLoggingStage)
            {
                return;
            }

            using (MessageBuffer mb = message.CreateBufferedCopy(this.MaxBufferSize))
            {
                using (Message copy = mb.CreateMessage())
                {
                    using (MemoryStream ms = new MemoryStream())
                    {
                        using (XmlTextWriter xtw = new XmlTextWriter(ms, Encoding.UTF8))
                        {
                            xtw.Formatting = Formatting.Indented;
                            copy.WriteMessage(xtw);
                            xtw.Flush();

                            ms.Seek(0, SeekOrigin.Begin);
                            using (StreamReader sr = new StreamReader(ms))
                            {
                                string xml = sr.ReadToEnd();
                                Trace.WriteLine(string.Format("Logging Message at stage {0}:{1}{2}",
                                    messageLoggingStage.ToString(),
                                    Environment.NewLine,
                                    xml));
                            }
                        }
                    }
                }

                message = mb.CreateMessage();
            }
        }
    }
}